#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Data extraction utility for physicochemical and absorption parameters
from Excel files without modifying the original file.
"""

import pandas as pd
import json
import os
from typing import Dict, Any, Optional
import logging

class ParameterExtractor:
    """Extract parameters from Excel files for PBPK modeling."""
    
    def __init__(self, excel_file_path: str):
        """Initialize with Excel file path."""
        self.excel_file_path = excel_file_path
        self.logger = logging.getLogger(__name__)
        
    def extract_physicochemical_parameters(self, sheet_name: str = "理化性质") -> Dict[str, Any]:
        """Extract physicochemical parameters from specified sheet."""
        try:
            df = pd.read_excel(self.excel_file_path, sheet_name=sheet_name)
            
            # Expected columns: 化合物名称, 分子量, logP, pKa, 溶解度, 渗透性
            parameters = {}
            
            for _, row in df.iterrows():
                compound_name = str(row.get('化合物名称', '')).strip()
                if compound_name:
                    parameters[compound_name] = {
                        'molecular_weight': float(row.get('分子量', 0)),
                        'logp': float(row.get('logP', 0)),
                        'pka': float(row.get('pKa', 7.0)),
                        'solubility': float(row.get('溶解度', 0)),
                        'permeability': float(row.get('渗透性', 0)),
                        'unit': {
                            'molecular_weight': 'g/mol',
                            'logp': 'unitless',
                            'pka': 'unitless',
                            'solubility': 'mg/mL',
                            'permeability': 'cm/s'
                        }
                    }
            
            self.logger.info(f"Extracted physicochemical parameters for {len(parameters)} compounds")
            return parameters
            
        except Exception as e:
            self.logger.error(f"Error extracting physicochemical parameters: {e}")
            return {}
    
    def extract_absorption_parameters(self, sheet_name: str = "吸收参数") -> Dict[str, Any]:
        """Extract absorption parameters from specified sheet."""
        try:
            df = pd.read_excel(self.excel_file_path, sheet_name=sheet_name)
            
            # Expected columns: 化合物名称, ka, F, Tlag, Tmax, Cmax
            parameters = {}
            
            for _, row in df.iterrows():
                compound_name = str(row.get('化合物名称', '')).strip()
                if compound_name:
                    parameters[compound_name] = {
                        'ka': float(row.get('ka', 0.1)),  # absorption rate constant
                        'bioavailability': float(row.get('F', 1.0)),  # fraction absorbed
                        'tlag': float(row.get('Tlag', 0)),  # lag time
                        'tmax': float(row.get('Tmax', 1.0)),  # time to max concentration
                        'cmax': float(row.get('Cmax', 0)),  # max concentration
                        'unit': {
                            'ka': '1/h',
                            'bioavailability': 'unitless',
                            'tlag': 'h',
                            'tmax': 'h',
                            'cmax': 'ng/mL'
                        }
                    }
            
            self.logger.info(f"Extracted absorption parameters for {len(parameters)} compounds")
            return parameters
            
        except Exception as e:
            self.logger.error(f"Error extracting absorption parameters: {e}")
            return {}
    
    def extract_all_parameters(self) -> Dict[str, Dict[str, Any]]:
        """Extract all parameters from the Excel file."""
        if not os.path.exists(self.excel_file_path):
            self.logger.error(f"Excel file not found: {self.excel_file_path}")
            return {}
        
        return {
            'physicochemical': self.extract_physicochemical_parameters(),
            'absorption': self.extract_absorption_parameters()
        }
    
    def save_extracted_data(self, output_dir: str = "data/extracted") -> str:
        """Save extracted data to JSON files without modifying original."""
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        all_params = self.extract_all_parameters()
        
        # Save combined data
        output_file = os.path.join(output_dir, "extracted_parameters.json")
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(all_params, f, ensure_ascii=False, indent=2)
        
        # Save individual parameter files
        for param_type, data in all_params.items():
            if data:
                individual_file = os.path.join(output_dir, f"{param_type}_parameters.json")
                with open(individual_file, 'w', encoding='utf-8') as f:
                    json.dump(data, f, ensure_ascii=False, indent=2)
        
        self.logger.info(f"Extracted data saved to {output_dir}")
        return output_dir
    
    def validate_parameters(self, parameters: Dict[str, Any]) -> Dict[str, list]:
        """Validate extracted parameters for completeness and consistency."""
        validation_results = {
            'missing_values': [],
            'invalid_ranges': [],
            'warnings': []
        }
        
        for compound, params in parameters.get('physicochemical', {}).items():
            # Check for missing critical values
            if params.get('molecular_weight', 0) <= 0:
                validation_results['missing_values'].append(f"{compound}: molecular_weight")
            
            if params.get('logp', 0) == 0:
                validation_results['warnings'].append(f"{compound}: logP is 0, check if intended")
        
        for compound, params in parameters.get('absorption', {}).items():
            if params.get('ka', 0) <= 0:
                validation_results['invalid_ranges'].append(f"{compound}: ka must be positive")
            
            if not 0 <= params.get('bioavailability', 0) <= 1:
                validation_results['invalid_ranges'].append(f"{compound}: bioavailability must be 0-1")
        
        return validation_results


def create_sample_data():
    """Create sample data structure for testing when Excel file is not available."""
    sample_data = {
        'physicochemical': {
            'acetaminophen': {
                'molecular_weight': 151.16,
                'logp': 0.46,
                'pka': 9.38,
                'solubility': 14.0,
                'permeability': 1.2e-05,
                'unit': {
                    'molecular_weight': 'g/mol',
                    'logp': 'unitless',
                    'pka': 'unitless',
                    'solubility': 'mg/mL',
                    'permeability': 'cm/s'
                }
            },
            'ibuprofen': {
                'molecular_weight': 206.28,
                'logp': 3.97,
                'pka': 4.91,
                'solubility': 0.021,
                'permeability': 2.8e-05,
                'unit': {
                    'molecular_weight': 'g/mol',
                    'logp': 'unitless',
                    'pka': 'unitless',
                    'solubility': 'mg/mL',
                    'permeability': 'cm/s'
                }
            }
        },
        'absorption': {
            'acetaminophen': {
                'ka': 2.5,
                'bioavailability': 0.88,
                'tlag': 0.1,
                'tmax': 0.75,
                'cmax': 12000,
                'unit': {
                    'ka': '1/h',
                    'bioavailability': 'unitless',
                    'tlag': 'h',
                    'tmax': 'h',
                    'cmax': 'ng/mL'
                }
            },
            'ibuprofen': {
                'ka': 1.8,
                'bioavailability': 0.95,
                'tlag': 0.2,
                'tmax': 1.5,
                'cmax': 8500,
                'unit': {
                    'ka': '1/h',
                    'bioavailability': 'unitless',
                    'tlag': 'h',
                    'tmax': 'h',
                    'cmax': 'ng/mL'
                }
            }
        }
    }
    
    # Save sample data
    output_dir = "data/sample"
    import os
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    with open(os.path.join(output_dir, "sample_parameters.json"), 'w', encoding='utf-8') as f:
        json.dump(sample_data, f, ensure_ascii=False, indent=2)
    
    return sample_data


if __name__ == "__main__":
    # Example usage
    excel_path = "理化性质、吸收参数表整理-0714.xlsx"
    
    if os.path.exists(excel_path):
        extractor = ParameterExtractor(excel_path)
        parameters = extractor.extract_all_parameters()
        validation = extractor.validate_parameters(parameters)
        output_dir = extractor.save_extracted_data()
        print(f"Data extracted and saved to {output_dir}")
    else:
        print(f"Excel file not found: {excel_path}")
        print("Creating sample data for testing...")
        sample_data = create_sample_data()
        print("Sample data created in data/sample/sample_parameters.json")